package com.ruoyi.web.controller.insect.controller;

import java.util.*;
import javax.servlet.http.HttpServletResponse;

import com.alibaba.fastjson.JSONObject;
import com.ruoyi.common.annotation.Log;
import com.ruoyi.common.utils.SecurityUtils;
import com.ruoyi.common.utils.uuid.IdUtils;
import com.ruoyi.insect.domain.*;
import com.ruoyi.insect.service.ICropDiseaseMappingService;
import com.ruoyi.insect.service.ITrainingService;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.security.access.prepost.PreAuthorize;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;
import com.ruoyi.common.core.controller.BaseController;
import com.ruoyi.common.core.domain.AjaxResult;
import com.ruoyi.common.enums.BusinessType;
import com.ruoyi.insect.service.IDiseaseLabelService;
import com.ruoyi.common.utils.poi.ExcelUtil;
import com.ruoyi.common.core.page.TableDataInfo;
import org.springframework.web.multipart.MultipartFile;

/**
 * 数据标注Controller
 */
@RestController
@RequestMapping("/insect/diseaselabel")
public class DiseaseLabelController extends BaseController {
    @Autowired
    private IDiseaseLabelService diseaseLabelService;

    @Autowired
    private ICropDiseaseMappingService cropCategoryService;

    private static final Logger log = LoggerFactory.getLogger(DiseaseLabelController.class);
    /**
     * 查询数据标注列表
     */
    @PreAuthorize("@ss.hasPermi('insect:diseaselabel:list')")
    @GetMapping("/list")
    public TableDataInfo list(DiseaseLabel diseaseLabel) {
        startPage();
        List<DiseaseLabel> list = diseaseLabelService.selectDiseaseLabelList(diseaseLabel);
        return getDataTable(list);
    }

    /**
     * 导出数据标注列表
     */
    @PreAuthorize("@ss.hasPermi('insect:diseaselabel:export')")
    @Log(title = "数据标注", businessType = BusinessType.EPORT)
    @PostMapping("/export")
    public void export(HttpServletResponse response, DiseaseLabel diseaseLabel) {
        List<DiseaseLabel> list = diseaseLabelService.selectDiseaseLabelList(diseaseLabel);
        ExcelUtil<DiseaseLabel> util = new ExcelUtil<DiseaseLabel>(DiseaseLabel.class);
        util.exportExcel(response, list, "数据标注数据");
    }

    /**
     * 获取数据标注详细信息
     */
    @PreAuthorize("@ss.hasPermi('insect:diseaselabel:query')")
    @GetMapping(value = "/{id}")
    public AjaxResult getInfo(@PathVariable("id") Long id) {
        return AjaxResult.success(diseaseLabelService.selectDiseaseLabelById(id));
    }

    /**
     * 新增数据标注
     */
    @PreAuthorize("@ss.hasPermi('insect:diseaselabel:add')")
    @Log(title = "数据标注", businessType = BusinessType.INSERT)
    @PostMapping
    public AjaxResult add(@RequestBody DiseaseLabel diseaseLabel) {
        return toAjax(diseaseLabelService.insertDiseaseLabel(diseaseLabel));
    }

    /**
     * 修改数据标注
     */
    @PreAuthorize("@ss.hasPermi('insect:diseaselabel:edit')")
    @Log(title = "数据标注", businessType = BusinessType.UPDATE)
    @PutMapping
    public AjaxResult edit(@RequestBody DiseaseLabel diseaseLabel) {
        return toAjax(diseaseLabelService.updateDiseaseLabel(diseaseLabel));
    }

    /**
     * 删除数据标注
     */
    @PreAuthorize("@ss.hasPermi('insect:diseaselabel:remove')")
    @Log(title = "数据标注", businessType = BusinessType.DELETE)
    @DeleteMapping("/{ids}")
    public AjaxResult remove(@PathVariable Long[] ids) {
        return toAjax(diseaseLabelService.deleteDiseaseLabelByIds(ids));
    }

    /**
     * 上传图片并处理GPS信息
     */
    @PreAuthorize("@ss.hasPermi('insect:diseaselabel:add')")
    @Log(title = "图片上传", businessType = BusinessType.INSERT)
    @PostMapping("/uploadImage")
    public AjaxResult uploadImage(
            @RequestParam("file") MultipartFile file,
            @RequestParam("deviceId") String deviceId,
            @RequestParam(value = "longitude", required = false) Double longitude,
            @RequestParam(value = "latitude", required = false) Double latitude) {
        try {
            if (file.isEmpty()) {
                return AjaxResult.error("上传文件不能为空");
            }

            // 获取当前用户名
            String userName = SecurityUtils.getUsername();

            // 处理文件上传
            String fileName = diseaseLabelService.processImageUpload(
                    file.getOriginalFilename(),
                    file.getBytes(),
                    deviceId,
                    userName,
                    longitude,
                    latitude
            );

            // 创建DiseaseLabel记录
            DiseaseLabel diseaseLabel = new DiseaseLabel();
            diseaseLabel.setOriginalImage(fileName);
            diseaseLabel.setDeviceId(deviceId);
            diseaseLabel.setGpsLongitude(longitude);
            diseaseLabel.setGpsLatitude(latitude);
            diseaseLabel.setPhotoTime(new java.util.Date());

            int result = diseaseLabelService.insertDiseaseLabel(diseaseLabel);

            if (result > 0) {
                // 返回包含文件名和ID的对象
                JSONObject responseData = new JSONObject();
                responseData.put("fileName", fileName);
                responseData.put("labelId", diseaseLabel.getId());
                return AjaxResult.success("上传成功", responseData);
            } else {
                return AjaxResult.error("保存记录失败");
            }

        } catch (Exception e) {
            logger.error("上传图片失败", e);
            return AjaxResult.error("上传失败: " + e.getMessage());
        }
    }

    /**
     * 保存标注数据
     */
    @PreAuthorize("@ss.hasPermi('insect:diseaselabel:edit')")
    @Log(title = "保存标注", businessType = BusinessType.UPDATE)
    @PostMapping("/saveAnnotations")
    public AjaxResult saveAnnotations(@RequestBody AnnotationRequest request) {
        try {
            DiseaseLabel diseaseLabel = diseaseLabelService.selectDiseaseLabelById(request.getLabelId());
            if (diseaseLabel == null) {
                return AjaxResult.error("标注记录不存在");
            }

            // 验证作物类型和病虫害名称
            if (request.getCropType() != null) {
                // 验证病虫害名称是否在允许的范围内
                List<String> allowedDiseases = cropDiseaseMappingService.getDiseaseListByCrop(request.getCropType());
                if (request.getAnnotations() != null) {
                    for (JSONObject annotation : request.getAnnotations()) {
                        String category = annotation.getString("category");
                        if (!allowedDiseases.contains(category)) {
                            return AjaxResult.error("病虫害名称 '" + category + "' 不在允许的范围内");
                        }
                    }
                }

                diseaseLabel.setCropType(request.getCropType());
            }

            // 构建标注数据结构
            JSONObject annotationData = new JSONObject();
            annotationData.put("file", diseaseLabel.getOriginalImage());
            annotationData.put("labels", request.getAnnotations());

            // 构建metadata
            JSONObject metadata = new JSONObject();
            metadata.put("lon", diseaseLabel.getGpsLongitude());
            metadata.put("lat", diseaseLabel.getGpsLatitude());
            metadata.put("device", diseaseLabel.getDeviceId());
            metadata.put("time", diseaseLabel.getPhotoTime());
            annotationData.put("metadata", metadata);

            String annotationDataStr = annotationData.toJSONString();
            diseaseLabel.setAnnotationData(annotationDataStr);

            logger.info("设置标注数据，长度: {}", annotationDataStr.length());

            // 如果有标注，设置第一个标注的病虫害名称
            if (request.getAnnotations() != null && !request.getAnnotations().isEmpty()) {
                JSONObject firstAnnotation = request.getAnnotations().get(0);
                diseaseLabel.setDiseaseName(firstAnnotation.getString("category"));
                logger.info("设置病虫害名称: {}", diseaseLabel.getDiseaseName());
            }

            // 使用新的方法保存，会生成标注图片
            logger.info("开始调用 updateDiseaseLabelWithAnnotation");
            int updateResult = diseaseLabelService.updateDiseaseLabelWithAnnotation(diseaseLabel);

            if (updateResult > 0) {
                // 重新查询获取更新后的数据，包括标注图片
                DiseaseLabel updatedLabel = diseaseLabelService.selectDiseaseLabelById(request.getLabelId());
                logger.info("标注保存成功，标注图片: {}", updatedLabel.getAnnotatedImage());

                // 返回包含标注图片信息的响应
                JSONObject responseData = new JSONObject();
                responseData.put("annotatedImage", updatedLabel.getAnnotatedImage());
                responseData.put("annotationCount", request.getAnnotations() != null ? request.getAnnotations().size() : 0);

                return AjaxResult.success("标注保存成功", responseData);
            } else {
                logger.error("标注保存失败，updateResult: {}", updateResult);
                return AjaxResult.error("标注保存失败");
            }

        } catch (Exception e) {
            logger.error("保存标注失败", e);
            return AjaxResult.error("保存失败: " + e.getMessage());
        }
    }

    // 在 DiseaseLabelController 中替换原有的作物服务引用
    @Autowired
    private ICropDiseaseMappingService cropDiseaseMappingService;

    @GetMapping("/getDiseasesByCrop")
    public AjaxResult getDiseasesByCrop(@RequestParam String cropType) {
        // 根据作物类型返回对应的病虫害列表
        Map<String, List<String>> cropDiseaseMap = new HashMap<>();
        cropDiseaseMap.put("水稻", Arrays.asList("稻瘟病", "稻飞虱", "稻纵卷叶螟", "纹枯病", "白叶枯病"));
        cropDiseaseMap.put("小麦", Arrays.asList("小麦锈病", "小麦白粉病", "赤霉病", "蚜虫", "吸浆虫"));
        cropDiseaseMap.put("玉米", Arrays.asList("玉米大斑病", "玉米小斑病", "玉米螟", "粘虫", "丝黑穗病"));
        cropDiseaseMap.put("大豆", Arrays.asList("大豆花叶病毒病", "大豆胞囊线虫病", "食心虫", "豆天蛾"));
        cropDiseaseMap.put("蔬菜", Arrays.asList("菜青虫", "小菜蛾", "霜霉病", "软腐病", "白粉病"));

        List<String> diseases = cropDiseaseMap.getOrDefault(cropType, new ArrayList<>());
        return AjaxResult.success(diseases);
    }

    // 在DiseaseLabelController类中添加以下方法

    @Autowired
    private ITrainingService trainingService;
    // ==================== 训练相关接口 ====================

    /**
     * 开始模型训练
     */
    @PostMapping("/startTraining")
    public AjaxResult startTraining(@RequestBody TrainForm trainForm) {
        try {
            log.info("开始模型训练，参数: {}", trainForm);

            // 创建训练任务
            TrainingTask task = new TrainingTask();
            task.setId(IdUtils.fastSimpleUUID());
            task.setTaskName(trainForm.getTaskName());
            task.setDescription(trainForm.getDescription());
            task.setEpochs(trainForm.getEpochs());
            task.setLearningRate(trainForm.getLearningRate());
            task.setBatchSize(trainForm.getBatchSize());
            task.setModelType(trainForm.getModelType());
            task.setUsePretrained(trainForm.getUsePretrained() ? 1 : 0);
            task.setStatus("running");
            task.setStartTime(new Date());
            task.setCreateBy(SecurityUtils.getUsername());

            // 异步执行训练
            trainingService.startTrainingAsync(task);

            return AjaxResult.success("训练任务已启动", task);
        } catch (Exception e) {
            log.error("开始训练失败", e);
            return AjaxResult.error("开始训练失败: " + e.getMessage());
        }
    }

    /**
     * 分割数据集
     */
    @PostMapping("/splitDataset")
    public AjaxResult splitDataset() {
        try {
            String result = trainingService.splitDataset();
            return AjaxResult.success("数据集分割成功", result);
        } catch (Exception e) {
            log.error("数据集分割失败", e);
            return AjaxResult.error("数据集分割失败: " + e.getMessage());
        }
    }

    /**
     * 停止训练
     */
    @PostMapping("/stopTraining/{taskId}")
    public AjaxResult stopTraining(@PathVariable String taskId) {
        try {
            trainingService.stopTraining(taskId);
            return AjaxResult.success("训练任务已停止");
        } catch (Exception e) {
            log.error("停止训练失败", e);
            return AjaxResult.error("停止训练失败: " + e.getMessage());
        }
    }

    /**
     * 获取训练状态
     */
    @GetMapping("/trainingStatus/{taskId}")
    public AjaxResult getTrainingStatus(@PathVariable String taskId) {
        try {
            TrainingTask task = trainingService.getTrainingStatus(taskId);
            return AjaxResult.success(task);
        } catch (Exception e) {
            log.error("获取训练状态失败", e);
            return AjaxResult.error("获取训练状态失败: " + e.getMessage());
        }
    }

    /**
     * 获取训练历史
     */
    @GetMapping("/trainingHistory")
    public AjaxResult getTrainingHistory() {
        try {
            List<TrainingTask> history = trainingService.getTrainingHistory();
            return AjaxResult.success(history);
        } catch (Exception e) {
            log.error("获取训练历史失败", e);
            return AjaxResult.error("获取训练历史失败: " + e.getMessage());
        }
    }

    /**
     * 获取训练日志
     */
    @GetMapping("/trainingLog/{taskId}")
    public AjaxResult getTrainingLog(@PathVariable String taskId) {
        try {
            String logContent = trainingService.getTrainingLog(taskId);
            return AjaxResult.success(logContent);
        } catch (Exception e) {
            log.error("获取训练日志失败", e);
            return AjaxResult.error("获取训练日志失败: " + e.getMessage());
        }
    }

    /**
     * 下载模型
     */
    @GetMapping("/downloadModel/{taskId}")
    public AjaxResult downloadModel(@PathVariable String taskId) {
        try {
            String modelPath = trainingService.getModelPath(taskId);
            return AjaxResult.success("模型路径: " + modelPath);
        } catch (Exception e) {
            log.error("获取模型路径失败", e);
            return AjaxResult.error("获取模型路径失败: " + e.getMessage());
        }
    }

    /**
     * 导出用户标注数据为COCO格式
     */
    @PostMapping("/exportUserAnnotations")
    public AjaxResult exportUserAnnotations() {
        try {
            String result = trainingService.exportUserAnnotationsToCoco();
            return AjaxResult.success("标注数据导出成功", result);
        } catch (Exception e) {
            log.error("导出标注数据失败", e);
            return AjaxResult.error("导出标注数据失败: " + e.getMessage());
        }
    }

    /**
     * 系统诊断
     */
    @GetMapping("/diagnoseSystem")
    public AjaxResult diagnoseSystem() {
        try {
            Map<String, Object> diagnoseResult = trainingService.diagnoseSystem();
            return AjaxResult.success(diagnoseResult);
        } catch (Exception e) {
            log.error("系统诊断失败", e);
            return AjaxResult.error("系统诊断失败: " + e.getMessage());
        }
    }

    /**
     * 调试训练环境
     */
    @GetMapping("/debugTrainingEnvironment")
    public AjaxResult debugTrainingEnvironment() {
        try {
            Map<String, Object> debugInfo = trainingService.debugTrainingEnvironment();
            return AjaxResult.success(debugInfo);
        } catch (Exception e) {
            log.error("调试训练环境失败", e);
            return AjaxResult.error("调试训练环境失败: " + e.getMessage());
        }
    }

    /**
     * 测试简单训练
     */
    @GetMapping("/testSimpleTraining")
    public AjaxResult testSimpleTraining() {
        try {
            Map<String, Object> result = trainingService.testSimpleTraining();
            return AjaxResult.success(result);
        } catch (Exception e) {
            log.error("测试训练失败", e);
            return AjaxResult.error("测试训练失败: " + e.getMessage());
        }
    }

}
